#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gmssl/sm2.h>
#include <gmssl/sm4.h>
#include <gmssl/rand.h>

typedef struct {
    uint8_t iv[SM4_BLOCK_SIZE];           // IV
    uint8_t ciphertext[1024];             // SM4加密的密文
    size_t ciphertext_len;                // 密文长度
    uint8_t encrypted_key[1024];          // SM2加密的SM4密钥
    size_t encrypted_key_len;             // 加密密钥长度
    uint8_t signature[1024];              // 数字签名
    size_t signature_len;                 // 签名长度
} DigitalEnvelope;

void print_hex(const char *label, const uint8_t *data, size_t len) {
    printf("%s (%zu bytes): ", label, len);
    for(size_t i = 0; i < len; i++) {
        printf("%02X", data[i]);
    }
    printf("\n");
}

int main() {
    // 1. 读取Alice的私钥
    SM2_KEY alice_private_key;
    FILE *priv_file = fopen("alice_private_key.bin", "rb");
    if (!priv_file) {
        perror("无法打开Alice私钥文件");
        return EXIT_FAILURE;
    }
    fread(&alice_private_key, sizeof(SM2_KEY), 1, priv_file);
    fclose(priv_file);
    
    // 2. 读取Bob的公钥
    SM2_KEY bob_public_key;
    FILE *pub_file = fopen("bob_public_key.bin", "rb");
    if (!pub_file) {
        perror("无法打开Bob公钥文件");
        return EXIT_FAILURE;
    }
    fread(&bob_public_key, sizeof(SM2_KEY), 1, pub_file);
    fclose(pub_file);
    
    // 3. 读取明文
    FILE *plain_file = fopen("plain.txt", "rb");
    if (!plain_file) {
        perror("无法打开明文文件");
        return EXIT_FAILURE;
    }
    fseek(plain_file, 0, SEEK_END);
    size_t plaintext_len = ftell(plain_file);
    fseek(plain_file, 0, SEEK_SET);
    
    uint8_t *plaintext = malloc(plaintext_len);
    fread(plaintext, 1, plaintext_len, plain_file);
    fclose(plain_file);
    
    printf("明文内容: ");
    fwrite(plaintext, 1, plaintext_len, stdout);
    printf("\n");
    
    // 4. 生成SM4密钥
    uint8_t sm4_key[SM4_KEY_SIZE];
    if (rand_bytes(sm4_key, SM4_KEY_SIZE) != 1) {
        fprintf(stderr, "SM4密钥生成失败\n");
        free(plaintext);
        return EXIT_FAILURE;
    }
    print_hex("生成的SM4密钥", sm4_key, SM4_KEY_SIZE);
    
    // 5. 生成随机IV
    uint8_t iv[SM4_BLOCK_SIZE];
    if (rand_bytes(iv, SM4_BLOCK_SIZE) != 1) {
        fprintf(stderr, "IV生成失败\n");
        free(plaintext);
        return EXIT_FAILURE;
    }
    print_hex("生成的IV", iv, SM4_BLOCK_SIZE);
    
    // 6. SM4加密明文
    SM4_KEY enc_key;
    sm4_set_encrypt_key(&enc_key, sm4_key);
    
    size_t max_ciphertext_len = plaintext_len + SM4_BLOCK_SIZE;
    uint8_t *ciphertext = malloc(max_ciphertext_len);
    size_t ciphertext_len;
    
    if (sm4_cbc_padding_encrypt(&enc_key, iv, plaintext, plaintext_len, ciphertext, &ciphertext_len) != 1) {
        fprintf(stderr, "SM4加密失败\n");
        free(plaintext);
        free(ciphertext);
        return EXIT_FAILURE;
    }
    
    // 7. SM2加密SM4密钥
    uint8_t encrypted_key[1024];
    size_t encrypted_key_len = sizeof(encrypted_key);
    
    if (sm2_encrypt(&bob_public_key, sm4_key, SM4_KEY_SIZE, encrypted_key, &encrypted_key_len) != 1) {
        fprintf(stderr, "SM2加密失败\n");
        free(plaintext);
        free(ciphertext);
        return EXIT_FAILURE;
    }
    
    // 8. SM2签名密文
    SM2_SIGN_CTX sign_ctx;
    uint8_t signature[1024];
    size_t signature_len = sizeof(signature);
    
    sm2_sign_init(&sign_ctx, &alice_private_key, SM2_DEFAULT_ID, SM2_DEFAULT_ID_LENGTH);
    sm2_sign_update(&sign_ctx, ciphertext, ciphertext_len);
    if (sm2_sign_finish(&sign_ctx, signature, &signature_len) != 1) {
        fprintf(stderr, "SM2签名失败\n");
        free(plaintext);
        free(ciphertext);
        return EXIT_FAILURE;
    }
    
    // 9. 保存数字信封
    DigitalEnvelope envelope;
    memcpy(envelope.iv, iv, SM4_BLOCK_SIZE);  // 保存IV
    memcpy(envelope.ciphertext, ciphertext, ciphertext_len);
    envelope.ciphertext_len = ciphertext_len;
    memcpy(envelope.encrypted_key, encrypted_key, encrypted_key_len);
    envelope.encrypted_key_len = encrypted_key_len;
    memcpy(envelope.signature, signature, signature_len);
    envelope.signature_len = signature_len;
    
    FILE *env_file = fopen("digital_envelope.bin", "wb");
    if (!env_file) {
        perror("无法打开数字信封文件");
        free(plaintext);
        free(ciphertext);
        return EXIT_FAILURE;
    }
    fwrite(&envelope, sizeof(DigitalEnvelope), 1, env_file);
    fclose(env_file);
    
    printf("数字信封生成成功\n");
    print_hex("IV", envelope.iv, SM4_BLOCK_SIZE);
    print_hex("SM4密文", envelope.ciphertext, envelope.ciphertext_len);
    print_hex("加密的SM4密钥", envelope.encrypted_key, envelope.encrypted_key_len);
    print_hex("数字签名", envelope.signature, envelope.signature_len);
    
    free(plaintext);
    free(ciphertext);
    return EXIT_SUCCESS;
}
